package com.gogant.spider.http;

import com.gogant.spider.core.Connector;
import com.gogant.spider.core.SpiderException;
import com.gogant.spider.js.Page;
import com.gogant.spider.js.PhantomPool;
import com.gogant.spider.js.PhantomProxy;
import com.gogant.spider.util.CharsetUtils;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.Proxy;
import java.net.URL;
import java.net.URLEncoder;
import java.nio.charset.Charset;
import java.util.Set;
import java.util.zip.DeflaterInputStream;
import java.util.zip.GZIPInputStream;

public class HttpConnector implements Connector<HttpSpider, HttpRequest, HttpResponse, HttpURLConnection> {
    private static final int LIMIT_BYTE_LENGTH = 10 * 1024 * 1024;
    private static int contentLimitBytes = LIMIT_BYTE_LENGTH;
    private static String acceptEncoding = "gzip, deflate";
    private static String userAgent = "Mozilla/5.0 (X11; Ubuntu; Linux i686; rv:26.0) Gecko/20100101 Firefox/26.0";

    static {
        // 禁用自动跳转
        HttpURLConnection.setFollowRedirects(false);
    }

    public static void setAcceptEncoding(String acceptEncoding) {
        HttpConnector.acceptEncoding = acceptEncoding;
    }

    public static String getAcceptEncoding() {
        return acceptEncoding;
    }

    public static void setUserAgent(String userAgent) {
        HttpConnector.userAgent = userAgent;
    }

    public static String getUserAgent() {
        return userAgent;
    }

    public static void setContentLimitBytes(int contentLimitBytes) {
        HttpConnector.contentLimitBytes = contentLimitBytes;
    }

    public static int getContentLimitBytes() {
        return contentLimitBytes;
    }

    @Override
    public HttpRequest createRequest(URL url) throws SpiderException {
        HttpRequest req = new HttpRequest(url);
        if (acceptEncoding != null) {
            req.putHeader("Accept-Encoding", acceptEncoding);
        }
        if (userAgent != null) {
            req.putHeader("User-Agent", userAgent);
        }
        return req;
    }

    @Override
    public HttpURLConnection openConnection(URL url, Proxy proxy) throws SpiderException {
        try {
            if (proxy == null) {
                return (HttpURLConnection) url.openConnection();
            } else {
                return (HttpURLConnection) url.openConnection(proxy);
            }
        } catch (Exception e) {
            throw new SpiderException("open url connection failed: " + e.getMessage(), e);
        }

    }

    private void sendRequestParameters(HttpURLConnection conn, HttpRequest req) throws IOException {
        Set<String> paramKeys = req.paramKeySet();
        if (paramKeys == null || paramKeys.size() == 0) {
            return;
        }
        if (req.getMethod() != null) {
            conn.setRequestMethod(req.getMethod().toString());
        }
        OutputStream output = null;
        try {
            output = conn.getOutputStream();
            StringBuilder sb = new StringBuilder();
            for (String key : paramKeys) {
                String value = req.getParam(key);
                if (value == null) {
                    value = "";
                }
                if (sb.length() > 0) {
                    sb.append("&");
                }
                sb.append(key).append("=").append(URLEncoder.encode(value, "UTF-8"));
            }
            output.write(sb.toString().getBytes());
            output.flush();
        } catch (IOException e) {
            throw e;
        } finally {
            if (output != null) {
                try {
                    output.close();
                } catch (Exception e) {
                }
            }
        }
    }

    private InputStream getDecompressStream(String type, InputStream input) throws IOException {
        if ("gzip".equalsIgnoreCase(type)) {
            input = new GZIPInputStream(input);
        } else if ("deflate".equalsIgnoreCase(type)) {
            input = new DeflaterInputStream(input);
        }
        return input;
    }

    private void fillContentBytes(HttpResponse resp, InputStream input, int limit) throws IOException {
        if (limit <= 0) {
            limit = Integer.MAX_VALUE;
        }
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        byte[] buffer = new byte[512];
        int readLength = buffer.length > limit ? limit : buffer.length;
        for (int len; (len = input.read(buffer, 0, readLength)) > 0; ) {
            output.write(buffer, 0, len);
            limit -= len;
            if (limit <= 0) {
                break;
            } else {
                readLength = buffer.length > limit ? limit : buffer.length;
            }
        }
        byte[] bytes = output.toByteArray();
        resp.
                setContentBytes(bytes);
        if (resp.getContentLength() != bytes.length) {
            resp.setContentLength(bytes.length);
        }
    }

    private String fillContentType(HttpResponse resp, InputStream input) {
        String contentType = resp.getContentType();
        if (contentType != null) {
            return contentType;
        }
        try {
            contentType = HttpURLConnection.guessContentTypeFromStream(input);
        } catch (Exception e) {
        }
        if (contentType == null) {
            contentType = "";
        }
        resp.setContentType(contentType);
        return contentType;
    }

    private void fillContentCharset(HttpResponse resp, String contentType, Charset defaultCharset) {
        String charset = null;
        int index = contentType.indexOf("charset=");
        if (index > 0) {
            int end = contentType.indexOf(index, ';');
            if (end > 0) {
                charset = contentType.substring(index, end);
            } else {
                charset = contentType.substring(index + 8);
            }
            charset = charset.trim();
            resp.setCharset(charset);
        } else if (contentType.startsWith("text")) {
            byte[] bytes = resp.getContentBytes();
            String str = new String(bytes, 0, bytes.length > 1024 ? 1024 : bytes.length);
            int idx = str.indexOf("charset=");
            if (idx > 0) {
                int endIdx = str.indexOf('>', idx);
                if (endIdx > idx) {
                    StringBuilder sb = new StringBuilder();
                    for (int i = idx + 9; i < endIdx; i++) {
                        char c = str.charAt(i);
                        if (c == '\'' || c == '"' || c == ' ') {
                            break;
                        } else {
                            sb.append(c);
                        }
                    }
                    try {
                        charset = Charset.forName(sb.toString()).name();
                    } catch (Exception e) {
                    }
                }
            }
            if (charset == null) {
                charset = CharsetUtils.detectCharset(bytes, 0, bytes.length > 5120 ? 5120 : bytes.length);
                if (charset == null) {
                    charset = defaultCharset == null ? "utf-8" : defaultCharset.name();
                }
            }
            resp.setCharset(charset);
        }
    }

    private void processJavascriptHtml(HttpResponse resp) throws Exception {
        String text = resp.getTextContent();
        if (text == null || text.length() == 0) {
            return;
        }
        PhantomProxy proxy = null;
        try {
            proxy = PhantomPool.getInstance().getProxy();
        } catch (Exception e) {
            e.printStackTrace();
            return;
        }
        if (proxy == null) {
            return;
        }
        try {
            Page page = proxy.loadPage(resp.getURL(), text, null, null);
            if (page != null) {
                resp.setContent(page.getContent());
            }
        } finally {
            proxy.close();
        }
    }

    public HttpResponse handleConnection(HttpSpider spider, HttpRequest req, HttpURLConnection conn) throws SpiderException {
        conn.setDoInput(true);
        conn.setDoOutput(true);
        conn.setUseCaches(false);
        HttpResponse resp = null;
        try {
            int retryTimes = spider.getReconnectAttempts();
            while (true) {
                try {
                    conn.connect();
                    break;
                } catch (Exception e) {
                    if (--retryTimes < 0) {
                        throw e;
                    }
                }
            }
            InputStream input = null;
            try {
                // send request parameters
                sendRequestParameters(conn, req);
                resp = new HttpResponse(req, conn);
                if (resp.getCode() != HttpURLConnection.HTTP_OK) {
                    return resp;
                }
                input = conn.getInputStream();

                String encoding = resp.getContentEncoding();
                if (encoding != null) {
                    input = getDecompressStream(encoding, input);
                }
                String contentType = fillContentType(resp, input);
                fillContentBytes(resp, input, contentLimitBytes);
                if (contentType != null) {
                    fillContentCharset(resp, contentType, spider.getDefaultContentCharset());
                }
                // process ajax html content
                if (req.isEnableJavascript() && contentType != null && contentType.startsWith("text/html")) {
                    processJavascriptHtml(resp);
                }
            } finally {
                if (input != null) {
                    try {
                        input.close();
                    } catch (Exception e) {
                    }
                }
            }
        } catch (Exception e) {
            resp = null;
            throw new SpiderException("process http connection failed: " + e.getMessage(), e);
        } finally {
            conn.disconnect();
        }
        return resp;
    }

    @Override
    public void destroy() {
    }

    @Override
    public String toString() {
        return "HttpConnector";
    }
}
